#!/usr/bin/env python3
"""
MAESS智能体系统 - MySQL数据库初始化脚本
创建必要的数据库和表结构
"""

import sys
import os
from loguru import logger

try:
    import pymysql
except ImportError:
    logger.error("pymysql 模块未安装，正在安装...")
    os.system("pip install pymysql")
    import pymysql

# MySQL连接配置
MYSQL_HOST = '192.168.64.100'
MYSQL_PORT = 3306
MYSQL_USER = 'root'
MYSQL_PASSWORD = 'root'
MYSQL_DATABASE = 'maess'

def create_database():
    """创建MAESS数据库"""
    try:
        # 连接到MySQL服务器（不指定数据库）
        conn = pymysql.connect(
            host=MYSQL_HOST,
            port=MYSQL_PORT,
            user=MYSQL_USER,
            password=MYSQL_PASSWORD,
            connect_timeout=5,
            charset='utf8mb4'
        )
        
        with conn.cursor() as cursor:
            # 检查数据库是否存在
            cursor.execute("SHOW DATABASES LIKE %s", (MYSQL_DATABASE,))
            if cursor.fetchone():
                logger.info(f"数据库 '{MYSQL_DATABASE}' 已存在")
                return True
            
            # 创建数据库
            cursor.execute(f"CREATE DATABASE {MYSQL_DATABASE} CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci")
            logger.success(f"✅ 数据库 '{MYSQL_DATABASE}' 创建成功")
            
        conn.commit()
        conn.close()
        return True
        
    except Exception as e:
        logger.error(f"❌ 数据库创建失败: {e}")
        return False

def create_tables():
    """创建必要的表结构"""
    try:
        # 连接到MAESS数据库
        conn = pymysql.connect(
            host=MYSQL_HOST,
            port=MYSQL_PORT,
            user=MYSQL_USER,
            password=MYSQL_PASSWORD,
            database=MYSQL_DATABASE,
            charset='utf8mb4'
        )
        
        with conn.cursor() as cursor:
            # 创建用户表
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS users (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    username VARCHAR(50) UNIQUE NOT NULL,
                    email VARCHAR(100) UNIQUE NOT NULL,
                    password_hash VARCHAR(255) NOT NULL,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
                    is_active BOOLEAN DEFAULT TRUE
                )
            """)
            
            # 创建任务表
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS tasks (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    user_id INT,
                    task_type VARCHAR(50) NOT NULL,
                    status VARCHAR(20) DEFAULT 'pending',
                    input_data TEXT,
                    output_data TEXT,
                    error_message TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
                    completed_at TIMESTAMP NULL,
                    FOREIGN KEY (user_id) REFERENCES users(id)
                )
            """)
            
            # 创建文档表
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS documents (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    user_id INT,
                    filename VARCHAR(255) NOT NULL,
                    file_path VARCHAR(500) NOT NULL,
                    file_size INT,
                    file_type VARCHAR(50),
                    content_hash VARCHAR(64),
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
                    FOREIGN KEY (user_id) REFERENCES users(id)
                )
            """)
            
            # 创建系统配置表
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS system_config (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    config_key VARCHAR(100) UNIQUE NOT NULL,
                    config_value TEXT,
                    description TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
                )
            """)
            
            # 插入默认配置
            cursor.execute("""
                INSERT IGNORE INTO system_config (config_key, config_value, description) VALUES
                ('app_name', 'MAESS智能体系统', '应用名称'),
                ('app_version', '1.0.0', '应用版本'),
                ('max_file_size', '10485760', '最大文件大小(10MB)'),
                ('supported_file_types', 'pdf,doc,docx,txt,md', '支持的文件类型')
            """)
            
            logger.success("✅ 数据表创建成功")
            
        conn.commit()
        conn.close()
        return True
        
    except Exception as e:
        logger.error(f"❌ 数据表创建失败: {e}")
        return False

def main():
    """主函数"""
    logger.info("开始初始化MAESS数据库...")
    logger.info(f"连接配置: {MYSQL_HOST}:{MYSQL_PORT}")
    
    # 创建数据库
    if not create_database():
        return 1
    
    # 创建数据表
    if not create_tables():
        return 1
    
    logger.success("\n🎉 数据库初始化完成！")
    logger.info("现在可以重新运行服务检查脚本了")
    return 0

if __name__ == "__main__":
    try:
        exit_code = main()
        sys.exit(exit_code)
    except Exception as e:
        logger.error(f"初始化失败: {e}")
        sys.exit(1)